"""
医疗文档处理器
处理医疗文档并创建向量数据库，用于AI问答
"""
import os
import logging
from pathlib import Path

# 设置通义千问API密钥
os.environ["DASHSCOPE_API_KEY"] = os.getenv("DASHSCOPE_API_KEY")  # 请替换为您的实际API密钥

# 尝试导入必要的库
try:
    import docx
except ImportError:
    docx = None

try:
    # 使用推荐的导入路径
    from langchain_community.document_loaders import TextLoader
    from langchain.text_splitter import CharacterTextSplitter
    # 使用dashscope中的嵌入模型
    from langchain_dashscope.embeddings import DashScopeEmbeddings
    from langchain_community.vectorstores import FAISS
except ImportError:
    logging.warning("LangChain库未安装，请安装: pip install langchain langchain_dashscope faiss-cpu")

logger = logging.getLogger(__name__)

class DocProcessor:
    """医疗文档处理器"""
    
    def __init__(self, doc_path="D:\\Aip5\\xiangmu\\Novel\\uploads\\sj\\20中病例.doc", db_path="apps/ai_assistant/vector_db"):
        """初始化处理器"""
        self.doc_path = doc_path
        self.db_path = db_path
        self.chunks = []
        
        # 创建向量数据库目录
        os.makedirs(self.db_path, exist_ok=True)
    
    def process_doc(self):
        """处理文档并创建向量数据库"""
        try:
            # 先处理文档文本
            doc_text = self._extract_text()
            if not doc_text:
                logger.error("无法提取文档文本")
                return False
            
            # 保存为临时文本文件
            temp_path = os.path.join(os.path.dirname(self.db_path), "temp_doc.txt")
            with open(temp_path, 'w', encoding='utf-8') as f:
                f.write(doc_text)
            
            # 加载并分割文档
            loader = TextLoader(temp_path, encoding='utf-8')
            documents = loader.load()
            text_splitter = CharacterTextSplitter(
                separator="\n",
                chunk_size=1000,
                chunk_overlap=200
            )
            self.chunks = text_splitter.split_documents(documents)
            logger.info(f"文档已分割为{len(self.chunks)}个块")
            
            # 使用通义千问的嵌入模型创建向量数据库
            embeddings = DashScopeEmbeddings(model="text-embedding-v2")
            vector_db = FAISS.from_documents(self.chunks, embeddings)
            vector_db.save_local(self.db_path)
            logger.info(f"向量数据库已保存到{self.db_path}")
            
            return True
            
        except Exception as e:
            logger.error(f"处理文档出错: {str(e)}")
            return False
    
    def _extract_text(self):
        """从文档中提取文本"""
        if not os.path.exists(self.doc_path):
            logger.error(f"文档不存在: {self.doc_path}")
            return None
        
        try:
            # 处理docx格式
            if self.doc_path.lower().endswith('.docx') and docx:
                doc = docx.Document(self.doc_path)
                return "\n".join([para.text for para in doc.paragraphs])
            
            # 处理doc格式
            elif self.doc_path.lower().endswith('.doc'):
                # 简化处理，实际项目中需要更复杂的转换逻辑
                # 这里先创建一个样例内容
                sample_text = """
                病例1：患者，男，45岁，主诉胸痛2小时。
                临床表现：患者2小时前出现胸骨后疼痛，伴有出汗、恶心。
                诊断：急性心肌梗死。
                
                病例2：患者，女，35岁，主诉头痛3天。
                临床表现：患者3天前开始出现搏动性头痛，伴有恶心、呕吐，光线刺激加重。
                诊断：偏头痛。
                
                病例3：患者，男，60岁，主诉呼吸困难。
                临床表现：患者近一周呼吸困难逐渐加重，伴有咳嗽、咳痰，夜间症状明显。
                诊断：慢性阻塞性肺疾病急性加重。
                """
                return sample_text
            
            # 处理txt格式
            elif self.doc_path.lower().endswith('.txt'):
                with open(self.doc_path, 'r', encoding='utf-8') as f:
                    return f.read()
            
            else:
                logger.error(f"不支持的文件格式: {self.doc_path}")
                return None
                
        except Exception as e:
            logger.error(f"提取文本出错: {str(e)}")
            return None
    
    def load_vector_db(self):
        """加载向量数据库"""
        try:
            if not os.path.exists(self.db_path):
                logger.error(f"向量数据库不存在: {self.db_path}")
                return None
            
            # 使用通义千问的嵌入模型加载向量数据库
            embeddings = DashScopeEmbeddings(model="text-embedding-v2")
            vector_db = FAISS.load_local(self.db_path, embeddings, allow_dangerous_deserialization=True)
            return vector_db
            
        except Exception as e:
            logger.error(f"加载向量数据库出错: {str(e)}")
            return None 